#include "PKEncrypter.h"
#include <botan/botan.h>
#include <botan/rsa.h>
#include <botan/pubkey.h>
#include <botan/pk_filts.h>
#include <botan/oids.h>
#include <botan/pem.h>
#include <botan/base64.h>
#include <assert.h>
#include <fstream>
#include <algorithm>

namespace kryptos { namespace cryptography {

PKEncrypter::PKEncrypter(const PublicKey& pubkey)
    : m_key(pubkey)
{
}

// Encrypter interface
int PKEncrypter::encryptFile(const std::string &filepath, const std::string &output)
{
    assert(false);
    return 0;
}

int PKEncrypter::encryptFile(const std::string &filepath, std::ostream &output)
{
    assert(false);
    return 0;
}

int PKEncrypter::encryptAndEncodeFile(const std::string &filepath, std::ostream &output)
{
    assert(false);
    return 0;
}

int PKEncrypter::encryptBuffer(const char* buffer, int len, std::vector<char> &output)
{
    // Recreate public key from bytes
    Botan::AutoSeeded_RNG rng;
    Botan::AlgorithmIdentifier algoId(Botan::OIDS::lookup("RSA"), Botan::AlgorithmIdentifier::USE_NULL_PARAM);
    Botan::MemoryVector<Botan::byte> bytes((const unsigned char*)m_key.key().constData(), m_key.key().size());

    Botan::RSA_PublicKey pubkey(algoId, bytes);

    // build PKEncrypter and encrypt buffer
    Botan::PK_Encryptor_EME encryptor(pubkey, "EME1(SHA-256)");
    Botan::SecureVector<Botan::byte> encbytes = encryptor.encrypt((const unsigned char*)buffer, len, rng);

    output.resize(encbytes.size());
    std::copy(encbytes.begin(), encbytes.end(), output.begin());

    return 0;
}

int PKEncrypter::encryptAndEncodeBuffer(const char* buffer, int len, std::vector<char> &output)
{
    encryptBuffer(buffer, len, output);
    size_t consumed = 0;
    Botan::base64_encode(output.data(), (const Botan::byte*)output.data(), output.size(), consumed, true);
    return 0;
}

int PKEncrypter::decryptFile(const std::string &filepath, const std::string &output)
{
    assert(false);
    return 0;
}

int PKEncrypter::decryptFile(std::istream &filepath, std::ostream &output)
{
    assert(false);
    return 0;
}

int PKEncrypter::decryptBuffer(const char* buffer, int len, std::vector<char> &output)
{
    // Recreate public key from bytes
    Botan::AutoSeeded_RNG rng;
    Botan::AlgorithmIdentifier algoId(Botan::OIDS::lookup("RSA"), Botan::AlgorithmIdentifier::USE_NULL_PARAM);
    Botan::MemoryVector<Botan::byte> bytes((const unsigned char*)m_key.key().constData(), m_key.key().size());

    Botan::RSA_PrivateKey privkey(algoId, bytes, rng);

    // build PKEncrypter and decrypt buffer
    Botan::PK_Decryptor_EME decryptor(privkey, "EME1(SHA-256)");
    Botan::SecureVector<Botan::byte> decbytes = decryptor.decrypt((const unsigned char*)buffer, len);
    output.resize(decbytes.size());
    std::copy(decbytes.begin(), decbytes.end(), output.begin());

    return 0;
}

}}
